package org.example.myleet.p501;

import org.example.myleet.Utils.TreeNode;

import java.util.*;

public class Solution {

    public int[] findMode(TreeNode root) {
        //key: node value, value: frequency
        Map<Integer, Integer> frequencies = new HashMap<>();
        if (null == root) {
            return new int[]{};
        }
        frequencies.put(root.val, 1);
        if (null != root.left) {
            frequencies = visitNode(frequencies, root.left);
        }
        if (null != root.right) {
            frequencies = visitNode(frequencies, root.right);
        }
        int peak = 0;
        List<Integer> result = new ArrayList<>();
        for (Map.Entry<Integer, Integer> entry : frequencies.entrySet()) {
            Integer nodeValue = entry.getKey();
            Integer frequency = entry.getValue();
            if (frequency > peak) {
                peak = frequency;
                result.clear();
                result.add(nodeValue);
            } else if (frequency == peak) {
                result.add(nodeValue);
            }
        }
        int[] r = new int[result.size()];
        for (int i = 0; i < result.size(); i++) {
            r[i] = result.get(i);
        }
        return r;
    }

    private Map<Integer, Integer> visitNode(Map<Integer, Integer> frequencies, TreeNode node) {
        if (null == node) {
            return frequencies;
        }
        Integer nodeValue = node.val;
        if (frequencies.containsKey(nodeValue)) {
            Integer newFrequency = frequencies.get(nodeValue) + 1;
            frequencies.put(nodeValue, newFrequency);
        } else {
            frequencies.put(nodeValue, 1);
        }
        if (null != node.left) {
            frequencies = visitNode(frequencies, node.left);
        }
        if (null != node.right) {
            frequencies = visitNode(frequencies, node.right);
        }
        return frequencies;
    }
}
